#!/usr/bin/env python

from ssummo.cmd_options import Options

import os.path
import sys

import Bio.SeqIO
from numpy import mean, std

## Some data for guessing sequence type from extension
FORMATS = { 'dat' : 'embl',
            'fas' : 'fasta',
            'fna' : 'fasta',
            'fsa' : 'fasta',
            'fasta': 'fasta',
        }

# Add some references to the sequence type iterator.
def update_formats(formats=None):
    if formats is None:
        formats = FORMATS
    for form in formats:
        Bio.SeqIO._FormatToIterator.update(
                { form : Bio.SeqIO._FormatToIterator[formats[form]] }
            )

update_formats()

def guess_format( file_name ):
    """
    Given a file name, will extract the file suffix, and try and guess the file
    type from that. Returns a format that can be accepted as an argument to
    SeqIO.parse( .... )
    """
    extension = file_name.rsplit( '.' , 1 )[1]
    if extension in Bio.SeqIO._FormatToIterator:
        return extension
    else:
        print('{0} is an unknown file extension.'.format(extension))
        print('Assuming fasta format')
        # Update the formats so above msg is only printed once.
        update_formats( {extension : 'fasta'} )
        return 'fasta'

class Lengths( ):
    def __init__(self):
        self.lens = []
        self.mean = 0.
        self.nBases = 0
        self.minimum = 1E99
        self.maximum = 0
        self.name = u''

    def __add__(self, other):
        end = Lengths()
        end.lens = self.lens + other.lens
        end.mean = (self.mean + other.mean) / 2.
        end.nBases = self.nBases + other.nBases
        end.minimum = min(self.minimum, other.minimum)
        end.maximum = max(self.maximum, other.maximum)
        return end

    def __str__(self):
        return str('{0} {1} {2} {3} {4} {5} {6} = {7}'
                    .format(self.name.ljust(28),
                            str(self.minimum).ljust(6),
                            str(self.maximum).ljust(7),
                            str( len(self.lens)).ljust(8),
                            '{0:.3f}'.format(self.getMean()).ljust(9),
                            '{0:.3f}'.format(self.getSTD()).ljust(10),
                            self.nBases,
                            bases(self.nBases)))

    def getLens(self, inhandle ,format=None):
        if format is None:
            format = guess_format(inhandle.name)
        for seq in Bio.SeqIO.parse(inhandle, format):
            # Just get the first sequence to check whether it's an alignment.
            break
        if '-' in seq.seq.tostring():
            # If the first sequence has '-' characters, assume this is an
            # alignment file, and we only want the actual residues.
            alignment = True
        else:
            alignment = False
        inhandle.seek( 0 )
        if alignment:
            for seq in Bio.SeqIO.parse(inhandle, format):
                seqLen = len(seq.seq.tostring().replace('-',''))
                self.lens.append(seqLen)
                self.minimum = min(seqLen, self.minimum)
                self.maximum = max(seqLen, self.maximum)
                self.nBases += seqLen
        else:
            for seq in Bio.SeqIO.parse(inhandle, format):
                seqLen = len(seq)
                self.lens.append(seqLen)
                self.minimum = min(seqLen, self.minimum)
                self.maximum = max(seqLen, self.maximum)
                self.nBases += seqLen
        return self.lens

    def getMean(self):
        if self.mean == 0.:
            self.mean = mean(self.lens)
        return self.mean

    def getSTD(self):
        return std( self.lens )

    def nBases(self):
        if self.nBases == 0:
            self.nBases = sum( self.nBases )
        return self.nBases

def get_stats(file_name, lengths=None):
    lengths = Lengths()
    with file( file_name, 'r' ) as inFile:
        lens = lengths.getLens(inFile)
    lengths.name = os.path.realpath(inFile.name).rsplit(os.path.sep, 1)[-1]
    nseqs = len(lens)
    mean = lengths.getMean()
    std = lengths.getSTD()
    return (nseqs, mean, std)

def bases(nBases):
    nBases = int(nBases)
    if nBases > 1E12:   # Tb
        nBases = float( nBases / 1E12 )
        return '{0:.6f} Tb'.format(nBases)
    elif nBases > 1E9:  # Gb
        return '{0:.6f} Gb'.format( float( nBases / 1E9 ) )
    elif nBases > 1E6:  # Mb
        return '{0:.3f} Mb'.format( float( nBases / 1E6 ) )
    elif nBases > 1E3:  # kb
        return '{0:.3f} kb'.format( float( nBases / 1E3 ) )
    else:
        return nBases, ' b'

def iterfiles(file_names, format=None):
    for file_name in file_names:
        lengths = Lengths( )
        with file( file_name , 'r' ) as inFile:
            lengths.getLens( inFile,format=format )
        lengths.name = os.path.realpath(file_name).rsplit(os.path.sep, 1)[-1]
        yield lengths

def get_headers():
    header = str('{0} {1} {2} {3} {4} {5} {6}\n'.format(
                'File'.ljust(28),'Min'.ljust(6),
                'Max'.ljust(7), 'nSeqs'.ljust(8),
                'mean'.ljust(9), 'Std.'.ljust(10),
                'nBases'.ljust(20)))
    header += str('{0} {1} {2} {3} {4} {5}'.format(
                'name'.ljust(28), 'len'.ljust(6),
                'len'.ljust(7), ' '.ljust(8),
                'length'.ljust(9), 'deviation'.ljust(10)))
    return header


def main(file_names, format=None):
    total_lengths = Lengths()
    print(get_headers())
    for lengths in iterfiles( file_names,format ): 
        print(lengths)
        total_lengths += lengths
    if len( file_names ) > 1:
        print('\n##  TOTALS  ##')
        print(total_lengths)

if __name__ == '__main__':

    options = Options()
    options.multiargs = ['-in']
    options.singleargs = ['-format']
    options.options.update( {'-in' : [] ,
                             '-format' : None } )
    options.parse_args(sys.argv[1:])
    file_names = options['-in']
    main( file_names, options['-format'] )
